from langchain_milvus import Milvus
from patagent.embedding.patsnap_embeddings import PatsnapEmbeddings
from langchain_openai import OpenAIEmbeddings
from patagent.tools import SwaggerParse
from patagent.constant import (
    TOPK_OPENAPI,
    MILVUS_URI, 
    MILVUS_TOKEN,
    MILVUS_DB,
    IS_REBUILD,
    PATSNAP_API_URL,
    PATSNAP_API_KEY,
    OPENAI_MODEL_HEADERS
)
from pymilvus import (
    Collection,
    CollectionSchema,
    DataType,
    FieldSchema,
    WeightedRanker,
    connections,
    utility
)
import json


embedding_func = OpenAIEmbeddings(openai_api_base=PATSNAP_API_URL, openai_api_key=PATSNAP_API_KEY, default_headers=OPENAI_MODEL_HEADERS)
connections.connect(uri=MILVUS_URI, token=MILVUS_TOKEN, db_name=MILVUS_DB)
col_name = "rag_openapi"
pk_field = "pk"
vector_field = "vector"
text_field = "text"
metadata_field = "metadata"

if IS_REBUILD:
    Collection(col_name).drop()

# Init or load collection
if not utility.has_collection(col_name):
    print('---INIT SWAGGER COLLECTION---')
    
    # Load
    sa = SwaggerParse()
    sa.main()
    docs_list = sa.api_name_list

    docs = []
    metadatas = []
    i = 0
    for doc in docs_list:
        docs.append(doc)
        if sa.api_method_list[i].upper() == 'GET':
            parameters = {
                k: v['default'] if 'default' in v else ''
                for k, v in json.loads(sa.api_params_path_list[i]).items()
            }
        else:
            parameters = json.loads(sa.api_params_body_list[i])
        metadatas.append({
            "name": doc,
            "url": sa.api_uri_list[i],
            "method": sa.api_method_list[i],
            "parameters": parameters,
        })
        i += 1
    dense_dim = len(embedding_func.embed_query(docs[1]))
    
    fields = [
        FieldSchema(
            name=pk_field,
            dtype=DataType.VARCHAR,
            is_primary=True,
            auto_id=True,
            max_length=100,
        ),
        FieldSchema(name=vector_field, dtype=DataType.FLOAT_VECTOR, dim=dense_dim),
        FieldSchema(name=text_field, dtype=DataType.VARCHAR, max_length=65535),
        FieldSchema(name=metadata_field, dtype=DataType.JSON),
    ]
    
    schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
    collection = Collection(
        name=col_name, schema=schema, consistency_level="Strong"
    )
    dense_index = {"index_type": "FLAT", "metric_type": "IP"}
    collection.create_index(vector_field, dense_index)
    collection.flush()

    entities = []
    i = 0
    for doc in docs:
        entity = {
            text_field: doc,
            vector_field: embedding_func.embed_query(doc),
            metadata_field: metadatas[i]
        }
        entities.append(entity)
        i += 1
    collection.insert(entities)
    collection.load()
else:
    collection = Collection(col_name)
    collection.load()
print('---LOADED SWAGGER COLLECTION---')

vector_store = Milvus(
    collection_name=col_name,
    embedding_function=embedding_func,
    vector_field=vector_field,
    connection_args={"uri": MILVUS_URI, "token": MILVUS_TOKEN}
)
api_retriever = vector_store.as_retriever(search_kwargs={'k': TOPK_OPENAPI})